#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
瞬变源分类程序
根据观测图像、模板图像和残差图像，对瞬变源候选体进行分类
"""

import numpy as np
import sep
from astropy.io import fits
import json
import os
from typing import List, Dict, Tuple


class ImageProperties:
    """图像属性类，用于计算和存储单个图像的属性"""
    
    def __init__(self, image: np.ndarray, x: int, y: int):
        """
        初始化图像属性
        
        参数:
            image: 完整图像 (2D numpy array)
            x, y: 目标坐标
        """
        self.image = image
        self.x = int(x)
        self.y = int(y)
        
        # 计算背景值和RMS噪声值
        self.background, self.rms = self._estimate_background_rms()
        
        # 裁剪7x7窗口
        self.cutout = self._cut_window(7)
        
        # 计算基本属性和峰值位置
        self.peak_flux, self.peak_x, self.peak_y = self._measure_peak_flux()
        self.flux3 = self._box_sum(3)
        self.flux5 = self._box_sum(5)
        self.fwhm = self._estimate_fwhm()
        self.sharpness = self._sharpness()
        
        # 计算信噪比(SNR)
        self.snr = self.peak_flux / (self.rms + 1e-6)
        
        # 计算轮廓比例
        self.peak_to_3x3 = self.peak_flux / (self.flux3 + 1e-6)
        self.flux3_to_5x5 = self.flux3 / (self.flux5 + 1e-6)
    
    def _estimate_background_rms(self) -> Tuple[float, float]:
        """
        估算图像的背景值和RMS噪声值
        
        返回:
            (background, rms) 元组
        """
        try:
            # 使用SEP估算背景和噪声
            bkg_obj = sep.Background(self.image.astype(np.float32))
            background = bkg_obj.globalback
            rms = bkg_obj.globalrms
        except:
            # 如果SEP失败，使用简单估算
            background = np.median(self.image)
            rms = np.std(self.image)
        return background, rms
    
    def _cut_window(self, size: int) -> np.ndarray:
        """
        裁剪以(x,y)为中心的size*size窗口
        
        参数:
            size: 窗口大小
        
        返回:
            裁剪后的图像窗口
        """
        half = size // 2
        x1, x2 = max(0, self.x - half), min(self.image.shape[1], self.x + half + 1)
        y1, y2 = max(0, self.y - half), min(self.image.shape[0], self.y + half + 1)
        cut = self.image[y1:y2, x1:x2]
        
        # 如果窗口小于指定大小，填充零
        if cut.shape != (size, size):
            padded = np.zeros((size, size))
            padded[:cut.shape[0], :cut.shape[1]] = cut
            return padded
        return cut
    
    def _measure_peak_flux(self) -> Tuple[float, int, int]:
        """
        测量峰值流量（扣除背景）和峰值位置坐标
        
        返回:
            (peak_flux, peak_x, peak_y) 元组
        """
        # 找到峰值在裁剪窗口中的位置
        max_idx = np.argmax(self.cutout)
        peak_y_local, peak_x_local = np.unravel_index(max_idx, self.cutout.shape)
        
        # 计算峰值在原始图像中的实际坐标
        half = self.cutout.shape[0] // 2
        peak_x = self.x + (peak_x_local - half)
        peak_y = self.y + (peak_y_local - half)
        
        # 计算峰值流量（扣除背景）
        peak_flux = self.cutout[peak_y_local, peak_x_local] - self.background
        
        return peak_flux, int(peak_x), int(peak_y)
    
    def _box_sum(self, size: int) -> float:
        """
        取中心的size*size方框内的总流量
        
        参数:
            size: 方框大小
        
        返回:
            方框内总流量（扣除背景）
        """
        if self.cutout.shape[0] < size or self.cutout.shape[1] < size:
            return 0
            
        c = size // 2
        cy, cx = self.cutout.shape[0] // 2, self.cutout.shape[1] // 2
        y1, y2 = cy - c, cy + c + 1
        x1, x2 = cx - c, cx + c + 1
        
        tcutout = self.cutout[y1:y2, x1:x2]
        tcutout[tcutout < 0] = 0
        box_flux = np.sum(tcutout)
        return box_flux - size * size * self.background
    
    def _estimate_fwhm(self) -> float:
        """
        粗略估算FWHM（半高全宽）
        
        返回:
            估算的FWHM值
        """
        peak = np.max(self.cutout)
        half_max = peak / 2
        yy, xx = np.where(self.cutout > half_max)
        
        if len(xx) == 0:
            return 0
        
        fwhm_x = xx.max() - xx.min() if len(xx) > 0 else 0
        fwhm_y = yy.max() - yy.min() if len(yy) > 0 else 0
        return max(fwhm_x, fwhm_y)
    
    def _sharpness(self) -> float:
        """
        判断是否非常陡峭（峰值流量与3x3流量比例）
        
        返回:
            陡峭度指标
        """
        flux3 = self.flux3
        peak = self.peak_flux
        
        if flux3 <= 0:
            return 0
        return peak / flux3


class Transient:
    """瞬变候选体类，用于存储和计算单个候选体的属性"""
    
    def __init__(self, obs_img: np.ndarray, tpl_img: np.ndarray, res_img: np.ndarray, 
                 x: int, y: int):
        """
        初始化瞬变候选体
        
        参数:
            obs_img: 观测图像 (2D numpy array)
            tpl_img: 模板图像 (2D numpy array)
            res_img: 残差图像 (2D numpy array)
            x, y: 目标坐标
        """
        self.x = int(x)
        self.y = int(y)
        
        # 创建各图像的属性对象
        self.obs_props = ImageProperties(obs_img, x, y)
        self.tpl_props = ImageProperties(tpl_img, x, y)
        self.res_props = ImageProperties(res_img, x, y)
        
        # 获取各图像的背景值和RMS噪声值
        self.obs_bkg = self.obs_props.background
        self.obs_rms = self.obs_props.rms
        self.tpl_bkg = self.tpl_props.background
        self.tpl_rms = self.tpl_props.rms
        self.res_bkg = self.res_props.background
        self.res_rms = self.res_props.rms
        
        # 获取各图像的属性
        self.peak_flux_obs = self.obs_props.peak_flux
        self.peak_x_obs = self.obs_props.peak_x
        self.peak_y_obs = self.obs_props.peak_y
        self.flux3_obs = self.obs_props.flux3
        self.flux5_obs = self.obs_props.flux5
        self.fwhm_obs = self.obs_props.fwhm
        self.sharpness_obs = self.obs_props.sharpness
        self.snr_obs = self.obs_props.snr
        
        self.peak_flux_tpl = self.tpl_props.peak_flux
        self.peak_x_tpl = self.tpl_props.peak_x
        self.peak_y_tpl = self.tpl_props.peak_y
        self.flux3_tpl = self.tpl_props.flux3
        self.flux5_tpl = self.tpl_props.flux5
        
        self.peak_flux_res = self.res_props.peak_flux
        self.peak_x_res = self.res_props.peak_x
        self.peak_y_res = self.res_props.peak_y
        self.flux3_res = self.res_props.flux3
        self.flux5_res = self.res_props.flux5
        self.fwhm_res = self.res_props.fwhm
        self.sharpness_res = self.res_props.sharpness
        
        # 计算轮廓比例
        self.peak_to_3x3_obs = self.obs_props.peak_to_3x3
        self.peak_to_3x3_tpl = self.tpl_props.peak_to_3x3
        self.peak_to_3x3_res = self.res_props.peak_to_3x3
        self.flux3_to_5x5_obs = self.obs_props.flux3_to_5x5
        self.flux3_to_5x5_tpl = self.tpl_props.flux3_to_5x5
        self.flux3_to_5x5_res = self.res_props.flux3_to_5x5
    
    
    
    def is_source(self) -> bool:
        """判断是否有源（信噪比大于5）"""
        return self.snr_obs > 5
    
    def to_json(self) -> Dict:
        """
        将瞬变候选体的属性转换为JSON格式的字典
        
        返回:
            包含所有属性的字典，按obs、tpl、res三幅图像组织
        """
        return {
            'x': int(self.x),
            'y': int(self.y),
            'obs': {
                'background': float(self.obs_bkg),
                'rms': float(self.obs_rms),
                'peak_flux': float(self.peak_flux_obs),
                'peak_x': int(self.peak_x_obs),
                'peak_y': int(self.peak_y_obs),
                'flux3': float(self.flux3_obs),
                'flux5': float(self.flux5_obs),
                'fwhm': float(self.fwhm_obs),
                'sharpness': float(self.sharpness_obs),
                'snr': float(self.snr_obs),
                'peak_to_3x3': float(self.peak_to_3x3_obs),
                'flux3_to_5x5': float(self.flux3_to_5x5_obs)
            },
            'tpl': {
                'background': float(self.tpl_bkg),
                'rms': float(self.tpl_rms),
                'peak_flux': float(self.peak_flux_tpl),
                'peak_x': int(self.peak_x_tpl),
                'peak_y': int(self.peak_y_tpl),
                'flux3': float(self.flux3_tpl),
                'flux5': float(self.flux5_tpl),
                'fwhm': float(self.tpl_props.fwhm),
                'sharpness': float(self.tpl_props.sharpness),
                'snr': float(self.tpl_props.snr),  
                'peak_to_3x3': float(self.peak_to_3x3_tpl),
                'flux3_to_5x5': float(self.flux3_to_5x5_tpl)
            },
            'res': {
                'background': float(self.res_bkg),
                'rms': float(self.res_rms),
                'peak_flux': float(self.peak_flux_res),
                'peak_x': int(self.peak_x_res),
                'peak_y': int(self.peak_y_res),
                'flux3': float(self.flux3_res),
                'flux5': float(self.flux5_res),
                'fwhm': float(self.fwhm_res),
                'sharpness': float(self.sharpness_res),
                'snr': float(self.res_props.snr),  
                'peak_to_3x3': float(self.peak_to_3x3_res),
                'flux3_to_5x5': float(self.flux3_to_5x5_res)
            }
        }


class TransientClassify:
    """瞬变源分类器主类"""
    
    def __init__(self, obs_img: np.ndarray, tpl_img: np.ndarray, res_img: np.ndarray, 
                 candidates: List[Tuple[int, int]]):
        """
        初始化分类器
        
        参数:
            obs_img: 观测图像
            tpl_img: 模板图像
            res_img: 残差图像
            candidates: 候选体坐标列表 [(x,y), ...]
        """
        self.obs_img = obs_img
        self.tpl_img = tpl_img
        self.res_img = res_img
        self.candidates = candidates
    
    def measure(self) -> List[Transient]:
        """
        对每个候选体进行测量
        
        返回:
            Transient对象列表
        """
        transients = []
        for x, y in self.candidates:
            try:
                transient = Transient(self.obs_img, self.tpl_img, self.res_img, x, y)
                transients.append(transient)
            except Exception as e:
                print(f"测量候选体({x},{y})时出错: {e}")
        return transients
    
    def classify(self, transient: Transient) -> str:
        """
        对单个瞬变候选体进行分类
        
        参数:
            transient: Transient对象
        
        返回:
            分类结果字符串（返回transient_type.json中定义的key）
        """
        
        # 4. 是否瞬变源
        if self._is_transient(transient):
            return "is_transient"
        
        # 1. 是否暗像素
        if self._is_dark_pixel(transient):
            return "is_dark_pixel"
        
        # 2. 是否是热像素
        if self._is_hot_pixel(transient):
            return "is_hot_pixel"
        
        # 3. 是否是亮星的无效残差
        if self._is_bright_star_residual(transient):
            return "is_bright_star_residual"
        
        # 1. 是否有源
        if not transient.is_source():
            return "other"
        
        # 5. 未知类型
        return "other"
    
    def _is_dark_pixel(self, transient: Transient) -> bool:
        """判断是否是暗像素"""
        if (transient.obs_props.flux5 <= 0 and transient.tpl_props.flux5 <= -100) \
            or (transient.tpl_props.flux3 <= 0 and transient.tpl_props.flux5 <=0):
            return True
        return False
    
    def _is_transient(self, transient: Transient) -> bool:
        """判断是否是瞬变源"""
        if transient.obs_props.flux3 >= 0 and transient.tpl_props.flux5 >= -150 and transient.tpl_props.flux5 <= 1000:
            if transient.fwhm_res >= 2 and transient.res_props.snr >= 3 and transient.obs_props.snr >= 4:
                return True
            
        return False
            
    def _is_hot_pixel(self, transient: Transient) -> bool:
        """判断是否是热像素"""
        if transient.obs_props.flux5 >= 100 and transient.tpl_props.flux5 >= -100 \
            and transient.obs_props.flux5 <= 5000 and transient.tpl_props.flux5 <= 5000:
            if transient.sharpness_res > 100 or transient.fwhm_res <= 2:
                return True
            
        return False
    
    def _is_bright_star_residual(self, transient: Transient) -> bool:
        """判断是否是亮星残差"""
        if transient.obs_props.flux5 >= 5000 and transient.tpl_props.flux5 >= 5000:
            return True
            
        return False

    def get_classification_criteria(self, transient: Transient) -> dict:
        """
        获取各分类函数使用的属性信息
        
        Args:
            transient: Transient对象
            
        Returns:
            包含各分类函数使用属性的字典
        """
        criteria = {
            'is_source': {
                'attributes': {
                    'peak_flux_obs': transient.peak_flux_obs,
                    'obs_rms': transient.obs_rms,
                    'snr': transient.snr_obs,
                    'peak_to_3x3': transient.peak_to_3x3_obs,
                    'flux3_to_5x5': transient.flux3_to_5x5_obs
                },
                'condition': f"peak_flux_obs > 5 * obs_rms\n即: {transient.peak_flux_obs:.2f} > 5 * {transient.obs_rms:.2f} = {5 * transient.obs_rms:.2f}"
            },
            'is_hot_pixel': {
                'attributes': {
                    'peak_flux_tpl': transient.peak_flux_tpl,
                    'tpl_rms': transient.tpl_rms,
                    'sharpness_obs': transient.sharpness_obs,
                    'sharpness_res': transient.sharpness_res,
                    'fwhm_obs': transient.fwhm_obs,
                    'snr': transient.snr_obs,
                    'peak_to_3x3': transient.peak_to_3x3_obs,
                    'flux3_to_5x5': transient.flux3_to_5x5_obs
                },
                'condition': (
                    f"条件1: peak_flux_tpl <= 3 * tpl_rms 且 sharpness_obs > 0.7 且 sharpness_res > 0.7 且 fwhm_obs <= 2\n"
                    f"即: {transient.peak_flux_tpl:.2f} <= 3 * {transient.tpl_rms:.2f} = {3 * transient.tpl_rms:.2f} 且 "
                    f"{transient.sharpness_obs:.2f} > 0.7 且 {transient.sharpness_res:.2f} > 0.7 且 {transient.fwhm_obs:.2f} <= 2\n"
                    f"或条件2: sharpness_obs > 0.8 且 sharpness_res > 0.8 且 fwhm_obs <= 1.5\n"
                    f"即: {transient.sharpness_obs:.2f} > 0.8 且 {transient.sharpness_res:.2f} > 0.8 且 {transient.fwhm_obs:.2f} <= 1.5"
                )
            },
            'is_bright_star_residual': {
                'attributes': {
                    'sharpness_obs': transient.sharpness_obs,
                    'peak_flux_tpl': transient.peak_flux_tpl,
                    'tpl_rms': transient.tpl_rms,
                    'fwhm_obs': transient.fwhm_obs,
                    'snr': transient.snr_obs,
                    'peak_to_3x3': transient.peak_to_3x3_obs,
                    'flux3_to_5x5': transient.flux3_to_5x5_obs
                },
                'condition': (
                    f"sharpness_obs < 0.6 且 peak_flux_tpl > 3 * tpl_rms 且 fwhm_obs > 2\n"
                    f"即: {transient.sharpness_obs:.2f} < 0.6 且 {transient.peak_flux_tpl:.2f} > 3 * {transient.tpl_rms:.2f} = {3 * transient.tpl_rms:.2f} 且 {transient.fwhm_obs:.2f} > 2"
                )
            },
            'is_transient': {
                'attributes': {
                    'peak_flux_obs': transient.peak_flux_obs,
                    'obs_rms': transient.obs_rms,
                    'peak_flux_tpl': transient.peak_flux_tpl,
                    'tpl_rms': transient.tpl_rms,
                    'fwhm_obs': transient.fwhm_obs,
                    'fwhm_res': transient.fwhm_res,
                    'sharpness_obs': transient.sharpness_obs,
                    'sharpness_res': transient.sharpness_res,
                    'snr': transient.snr_obs,
                    'peak_to_3x3': transient.peak_to_3x3_obs,
                    'flux3_to_5x5': transient.flux3_to_5x5_obs
                },
                'condition': (
                    f"peak_flux_obs > 5 * obs_rms 且 peak_flux_tpl <= 3 * tpl_rms\n"
                    f"即: {transient.peak_flux_obs:.2f} > 5 * {transient.obs_rms:.2f} = {5 * transient.obs_rms:.2f} 且 "
                    f"{transient.peak_flux_tpl:.2f} <= 3 * {transient.tpl_rms:.2f} = {3 * transient.tpl_rms:.2f}\n"
                    f"且 |fwhm_obs - fwhm_res| < 2 且 |sharpness_obs - sharpness_res| < 0.3\n"
                    f"即: |{transient.fwhm_obs:.2f} - {transient.fwhm_res:.2f}| = {abs(transient.fwhm_obs - transient.fwhm_res):.2f} < 2 且 "
                    f"|{transient.sharpness_obs:.2f} - {transient.sharpness_res:.2f}| = {abs(transient.sharpness_obs - transient.sharpness_res):.2f} < 0.3"
                )
            }
        }
        return criteria


def classifyImage(fits_file: str, candidate_coords: List[Tuple[int, int]] = None) -> Dict:
    """
    分类单个FITS图像中的瞬变源
    
    参数:
        fits_file: FITS文件路径
        candidate_coords: 候选体坐标列表，默认为[(50,50)]
    
    返回:
        分类结果字典
    """
    if candidate_coords is None:
        candidate_coords = [(50, 50)]
    
    try:
        # 读取FITS文件
        with fits.open(fits_file) as hdul:
            # 解析三个图像
            obs_img = None
            tpl_img = None
            res_img = None
            
            for hdu in hdul:
                if 'IMAGE_TYPE' in hdu.header:
                    img_type = hdu.header['IMAGE_TYPE'].strip().upper()
                    if img_type == 'OBJT':
                        obs_img = hdu.data
                    elif img_type == 'TEMP':
                        tpl_img = hdu.data
                    elif img_type == 'DIFF':
                        res_img = hdu.data
            
            if obs_img is None or tpl_img is None or res_img is None:
                raise ValueError("FITS文件中缺少必要的图像类型")
            
            # 创建分类器
            classifier = TransientClassify(obs_img, tpl_img, res_img, candidate_coords)
            
            # 测量候选体
            transients = classifier.measure()
            
            # 分类结果
            results = []
            for transient in transients:
                classification = classifier.classify(transient)
                # 使用Transient类的to_json方法获取所有属性
                result_dict = transient.to_json()
                result_dict['classification'] = classification
                results.append(result_dict)
            
            return {
                'fits_file': fits_file,
                'results': results
            }
            
    except Exception as e:
        return {
            'fits_file': fits_file,
            'error': str(e),
            'results': []
        }


def batchClassify(transient_list_file: str, images_dir: str, output_file: str = None):
    """
    批量分类瞬变源
    
    参数:
        transient_list_file: 瞬变源列表JSON文件路径
        images_dir: 图像文件目录
        output_file: 输出文件路径，默认为transient_classify_rst.json
    """
    if output_file is None:
        # 使用瞬变源列表文件所在目录作为输出目录
        output_dir = os.path.dirname(transient_list_file)
        output_file = os.path.join(output_dir, 'transient_classify_rst.json')
    
    # 读取瞬变源列表
    with open(transient_list_file, 'r') as f:
        data = json.load(f)
    
    objects = data.get('objects', [])
    obj_number = len(objects)
    
    results = {
        'objects': [],
        'obj_number': obj_number
    }
    
    # 对每个瞬变源进行分类
    for obj in objects:
        name = obj.get('name', '')
        fits_filename = f"{name}_001.fit"
        fits_path = os.path.join(images_dir, fits_filename)
        
        if not os.path.exists(fits_path):
            print(f"警告: 图像文件不存在: {fits_path}")
            continue
        else:
            # 获取候选体坐标
            basic_info = obj.get('basic_info', {})
            # x = basic_info.get('x', 50)  # 默认为中心坐标
            # y = basic_info.get('y', 50)
            x = 50
            y = 50
            
            # 分类单个图像
            classification_result = classifyImage(fits_path, [(x, y)])
            
            if 'error' in classification_result:
                result = {
                    'name': name,
                    'classification': '未知类型',
                    'image_path': fits_path,
                    'error': classification_result['error']
                }
            else:
                # 取第一个候选体的分类结果
                if classification_result['results']:
                    first_result = classification_result['results'][0]
                    result = first_result
                    result['name'] = name
                    result['image_path'] = fits_path
                else:
                    result = {
                        'name': name,
                        'classification': '未知类型',
                        'image_path': fits_path,
                        'error': '没有检测到候选体'
                    }
        
        results['objects'].append(result)
        print(f"处理完成: {name} -> {result['classification']}")
    
    # 保存结果
    with open(output_file, 'w') as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    
    print(f"批量分类完成，结果保存在: {output_file}")
    return results


if __name__ == "__main__":
    # 处理指定路径的数据
    import sys
    
    date_str = "20251016_045220" #20251004_161049 20251007
    # 默认路径
    transient_list_file = f'data/datasets/{date_str}/transient_list.json'
    images_dir = f'data/datasets/{date_str}/images'
    output_file = f'data/datasets/{date_str}/transient_classify_rst.json'
    
    # 如果提供了命令行参数，则使用参数指定的路径
    if len(sys.argv) == 4:
        transient_list_file = sys.argv[1]
        images_dir = sys.argv[2]
        output_file = sys.argv[3]
    elif len(sys.argv) != 1:
        print("用法: python transient_classifier.py [瞬变源列表文件] [图像目录] [输出文件]")
        print("示例: python transient_classifier.py data/datasets/20251004_161049/transient_list.json data/datasets/20251004_161049/images data/datasets/20251004_161049/transient_classify_rst.json")
        sys.exit(1)
    
    print(f"开始处理瞬变源分类任务...")
    print(f"瞬变源列表文件: {transient_list_file}")
    print(f"图像目录: {images_dir}")
    print(f"输出文件: {output_file}")
    
    # 执行批量分类
    batchClassify(transient_list_file, images_dir, output_file)